import av
import numpy as np
from flask import Flask, Response
import cv2

from logger import log
from settings import settings

app = Flask(__name__)


# rtsp://admin:admin@192.168.110.100:8557/h264
def getRtspUrl(rtsp_code):
    # 遍历列表，寻找匹配的id
    for item in settings.streams:
        if item['id'] == int(rtsp_code.split("_")[-1]):
            # 如果找到匹配的id，获取对应的rtsp
            rtsp_url = item['rtsp']
            return rtsp_url
    else:
        log.info('No matching id found in the list.')


def generate_frames_av(rtsp_url):
    frame_count = 0  # 帧计数器
    try:
        container = av.open(rtsp_url,
                            format='rtsp',
                            options={'rtsp_transport': 'tcp',
                                     'framerate': '20',
                                     'max_delay': '5000000',
                                     'video_size': '640x360'},
                            timeout=5)
    except Exception as e:
        log.warn(f"generate_frames_av Failed to open RTSP stream: {e}")
        # 创建一个固定图片
        fixed_image = np.zeros((360, 640, 3), dtype=np.uint8)  # 黑色图片
        cv2.putText(fixed_image, 'Failed to open stream', (10, 180),
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
        ret, buffer = cv2.imencode('.jpg', fixed_image)
        if ret:
            frame = buffer.tobytes()
            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')
        return None
    for frame in container.decode(video=0):
        frame_count += 1
        # 每读取到两帧只处理一帧，去掉一半的帧
        if frame_count % 5 == 0:
            # 将帧转换为numpy数组
            img = frame.to_ndarray(format='bgr24')
            # 将帧调整为640x360
            img = cv2.resize(img, (640, 360))
            # 将帧编码为JPEG格式
            ret, buffer = cv2.imencode('.jpg', img)
            if not ret:
                continue
            frame = buffer.tobytes()
            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')


def generate_frames(rtsp_url):
    cap = cv2.VideoCapture(rtsp_url)
    if not cap.isOpened():
        log.info("generate_frames not cap.isOpened()")
        return None
    frame_count = 0  # 帧计数器
    while True:
        success, frame = cap.read()  # 从RTSP流读取帧
        if not success:
            log.error(f"取流失败:{rtsp_url}")
            break
        else:
            frame_count += 1
            # 每读取到两帧只处理一帧，去掉一半的帧
            if frame_count % 5 == 0:
                # 将帧调整为640x360
                frame = cv2.resize(frame, (640, 360))

                # 将帧编码为JPEG格式
                ret, buffer = cv2.imencode('.jpg', frame)
                frame = buffer.tobytes()
                # count += 1
                # if count > 100:
                #     break
                # 使用生成器返回帧
                yield (b'--frame\r\n'
                       b'Content-Type: image/jpg\r\n\r\n' + frame + b'\r\n')


@app.route('/video/<code>')
def video_feed(code):
    rtsp_url = getRtspUrl(code)
    log.info(rtsp_url)
    return Response(generate_frames_av(rtsp_url),
                    mimetype='multipart/x-mixed-replace; boundary=frame')


@app.route('/')
def hello():
    links = []
    for item in settings.streams:
        stream_id = item['id']
        links.append(f'<li><a href="/video/rtsp_{stream_id}">{item["name"]}</a></li>')
    links_html = '<ul>' + ''.join(links) + '</ul>'
    return f'<html><body><h1>Streams</h1>{links_html}</body></html>'


def start_http_server():
    log.info('开启http服务')
    app.run(host='127.0.0.1', port=5100)


if __name__ == '__main__':
    settings.init(code='lipeng_test', host='platform.wechain360.com')
    start_http_server()
